// kdtree.cpp
// 04/17/2011: created

#include "kdtree.h"
#include <QtCore/QVector>
#include "Log.h"

void KDTree::random_permutation(std::vector<int>& p)
{
  int n = p.size();
  for (int i = 0; i < n; i++) p[i] = i;
  for (int i = 0; i < n-1; i++) {
    int swap = rand() % (n-i) + i;
    int tmp = p[i]; p[i] = p[swap]; p[swap] = tmp;
  }
}

KDTree::KDTree(dataset_type& data)
  : Tree(data), m_box(n_dim()), m_dimList(*(new vint_type(n_dim())))
{
  random_permutation(m_dimList);
}

KDTree::KDTree(KDTree* parent, int s, int e)
  : Tree(parent, s, e), m_box(n_dim()), m_dimList(parent->m_dimList)
{
}

KDTree::~KDTree()
{
  if (!m_parent) {
    delete &m_dimList;
  }
}

int KDTree::split(int dim)
{
  // choose split value in dimension dim
  QVector<double> data(n());
  for_all_node_data(i) data[i] = get(i)[dim];
  qSort(data.begin(), data.end());

  double splitVal = data[n()/2];
  if (n()/2+1 < n()) splitVal = (data[n()/2] + data[n()/2+1]) / 2;

  // try splitting the node by moving indices with value < splitVal to the left
  int n_left = 0;
  for_all_node_data(i) if (get(i)[dim] < splitVal) {
    int tmp = m_idx[m_s+n_left]; m_idx[m_s+n_left] = m_idx[m_s+i]; m_idx[m_s+i] = tmp;
    n_left++;
  }
  if (n_left > 0 && n_left < n())
    return n_left;  // success, return number of data points in left node
  else
    return -1; // failure
}

void KDTree::buildTree()
{
  int n_left = -1;

  if (n() > 10) {  // only split big nodes
    int dimIdx = rand() % n_dim();
    //Log::m_log << "split node n=" << n() << endl;
    for (int dim = 0; dim < n_dim(); dim++, dimIdx = (dimIdx+1)%n_dim()) {
      n_left = split(m_dimList[dimIdx]);
      //Log::m_log << "try split at dim=" << m_dimList[dimIdx] << " n_left=" << n_left << endl;
      if (n_left != -1) {
        //Log::m_log << "split at dim=" << m_dimList[dimIdx] << endl;
        break;
      }
    }
  }

  if (n_left != -1) {
    // split
    KDTree *left = new KDTree(this, m_s, m_s+n_left);
    KDTree *right = new KDTree(this, m_s+n_left, m_s+n());
    m_child.push_back(left);
    m_child.push_back(right);
    left->buildTree();
    right->buildTree();
    // build node stats
    m_box.accumulate(left->m_box);
    m_box.accumulate(right->m_box);
  }
  else { // build leaf stats
    for_all_node_data(i) {
      m_box.accumulate(get(i));
    }
  }
}

int KDTree::findNN(const vec& query, double& minBound)
{
  double minDist = minBound;
  int minIdx = -1;
  double lowerBound = m_box.lowerBound(query);
  //Log::m_log << "lowerBound=" << lowerBound << endl;
  if (lowerBound >= minBound) {               // pruned this node
    //Log::m_log << "pruned" << endl;
    return minIdx;
  }
  if (m_child.size() == 0) {                 // the base case at leaf node
    for_all_node_data(i) {
      double dist = arma::norm(query-get(i), 2);
      //Log::m_log << "idx=" << originalIdx(i) << " dist=" << dist << endl;
      if (dist < minDist) {
        minDist = dist;
        minIdx = m_s+i;
        //Log::m_log << "better candidate minIdx=" << minIdx << endl;
      }
    }
  }
  else {
    std::vector<KDTree*> childs(2);
    childs[0] = static_cast<KDTree*>(m_child[0]);
    childs[1] = static_cast<KDTree*>(m_child[1]);
    if (childs[0]->m_box.lowerBound(query) > childs[1]->m_box.lowerBound(query)) {
      childs[0] = static_cast<KDTree*>(m_child[1]);
      childs[1] = static_cast<KDTree*>(m_child[0]);
      //Log::m_log << "change node order" << endl;
    }
    for (unsigned int i = 0; i < m_child.size(); i++) {
      KDTree* child = childs[i];
      int idx = child->findNN(query, minDist);
      if (idx != -1) minIdx = idx;   // find a better point dist < minDist
    }
  }
  minBound = minDist;
  return minIdx;
}

void KDTree::findKNN(const vec& query, std::vector<int>& knnIdx, std::vector<double>& minBound)
{
  double lowerBound = m_box.lowerBound(query);
  //Log::m_log << "lowerBound=" << lowerBound << endl;
  if (lowerBound >= minBound.back()) {               // pruned this node
    //Log::m_log << "pruned" << endl;
    return;
  }

  if (m_child.size() == 0) {
    Tree::findKNN(query, knnIdx, minBound);
  }
  else {
    std::vector<KDTree*> childs(2);
    childs[0] = static_cast<KDTree*>(m_child[0]);
    childs[1] = static_cast<KDTree*>(m_child[1]);
    if (childs[0]->m_box.lowerBound(query) > childs[1]->m_box.lowerBound(query)) {
      childs[0] = static_cast<KDTree*>(m_child[1]);
      childs[1] = static_cast<KDTree*>(m_child[0]);
      //Log::m_log << "change node order" << endl;
    }
    for (unsigned int i = 0; i < m_child.size(); i++) {
      KDTree* child = childs[i];
      child->findKNN(query, knnIdx, minBound);
    }
  }
}

// ================ BOUNDINGBOX ================
BoundingBox::BoundingBox(int dim) : m_min(dim), m_max(dim), n_dim(dim) 
{
  m_min.fill(POSITIVE_INFINITY);
  m_max.fill(NEGATIVE_INFINITY);
}

void BoundingBox::accumulate(const vec& v) 
{
  assert(v.n_elem == n_dim);
  for (int i = 0; i < n_dim; i++) {
    if (m_min[i] > v[i]) m_min[i] = v[i];
    if (m_max[i] < v[i]) m_max[i] = v[i];
  }
}

void BoundingBox::accumulate(const BoundingBox& b) 
{
  assert(b.n_dim == n_dim);
  for (int i = 0; i < n_dim; i++) {
    if (m_min[i] > b.m_min[i]) m_min[i] = b.m_min[i];
    if (m_max[i] < b.m_max[i]) m_max[i] = b.m_max[i];
  }
}

const vec& BoundingBox::min() const 
{ 
  return m_min; 
}

const vec& BoundingBox::max() const 
{ 
  return m_max; 
}

double BoundingBox::lowerBound(const vec& v) const
{
  double s = 0;
  for (int i = 0; i < n_dim; i++) {
    if (v[i] < m_min[i]) s += (v[i] - m_min[i]) * (v[i] - m_min[i]);
    if (v[i] > m_max[i]) s += (v[i] - m_max[i]) * (v[i] - m_max[i]);
  }
  return sqrt(s);
}
